package com.bwie.ai.websocket;

import org.springframework.stereotype.Component;
import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint("/msgServer/{userId}")
public class WebSocketServer {

    private static ConcurrentHashMap<String, Session> webSocketMap = new ConcurrentHashMap<>();
    private String userId;

    @OnOpen
    public void onOpen(Session session, @javax.websocket.server.PathParam("userId") String userId) {
        this.userId = userId;
        webSocketMap.put(userId, session);
        System.out.println("用户 " + userId + " 连接");
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        System.out.println("收到消息: " + message);
    }

    @OnError
    public void onError(Session session, Throwable error) {
        System.out.println("用户 " + userId + " 连接异常: " + error.getMessage());
        error.printStackTrace();
        try {
            if (session.isOpen()) {
                session.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @OnClose
    public void onClose() {
        if (userId != null && !userId.isEmpty()) {
            webSocketMap.remove(userId);
            System.out.println("用户 " + userId + " 断开连接");
        }
    }
} 